package org.pombo.gtb;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import org.pombo.gtm.Action;
import org.pombo.gtm.Strategy;

public class StrategyBuilder extends Builder<Strategy> {

	@Override
	public void clear() {
		this.entity = mock(Strategy.class);
	}

	public StrategyBuilder withProbability(Action action, double probability) {
		when(this.entity.getProbabilityForAction(action)).thenReturn(probability);
		return this;
	}
}
